from torch import (
    tensor,
    no_grad,
    load
)
from core import MovieBiLSTM_CRF
from utils import conf, data

_, word2id = data.build_vocab()
id2tag = {v: k for k, v in data.labels.items()}
Entity = dict[str, str]

def __extract_entity(chars: list[str], tags: list[str]) -> Entity:
    # Initialize entities
    entities, curr_entities = [], []
    curr_type = None
    # Extract entities
    for char, tag in zip(chars, tags):
        if tag.startswith('B-'):
            if curr_entities:
                entities.append((
                    curr_type,
                    ''.join(curr_entities)
                ))
                curr_entities = []
            curr_type = tag.split('-')[1]
            curr_entities.append(char)
        elif tag.startswith('I-') and curr_entities:
            curr_entities.append(char)
        else:
            if curr_entities:
                entities.append((
                    curr_type,
                    ''.join(curr_entities)
                ))
                curr_entities = []
                curr_type = None
    # Add last entity
    if curr_entities:
        entities.append((
            curr_type,
            ''.join(curr_entities)
        ))
    # Return entities
    return {entity: entity_type for entity_type, entity in entities}

def predict(input_text: dict[str, str]) -> Entity:
    # Get input text
    text = input_text['text']
    x = [word2id.get(char, word2id.get('UNK', 0)) for char in text]
    x_test = tensor([x]).to(conf.device)
    mask = (x_test != 0).long()
    # Initialize model
    model = MovieBiLSTM_CRF()
    # Load model
    model.load_state_dict(load(conf.model.path_torch))
    model = model.to(conf.device)
    model.eval()
    # Predict
    with no_grad():
        predict = model(x_test, mask)
        tags = [id2tag[i] for i in predict[0]]
        # Extract entities
        chars = list(text)
        return __extract_entity(chars, tags)
